package org.apache.flink.runtime.metrics.groups;

import static java.lang.Thread.holdsLock;
import static org.apache.flink.util.Preconditions.checkNotNull;
import static org.apache.flink.util.Preconditions.checkState;

import java.util.HashMap;
import java.util.Map;

import javax.annotation.Nullable;

import org.apache.flink.annotation.Internal;
import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.runtime.metrics.MetricRegistry;
import org.apache.flink.runtime.resourceconsumer.ConsumeAttemptID;

/**
 * Special {@link org.apache.flink.metrics.MetricGroup} representing everything belonging to a
 * specific consumer, running on the ResourceProvider.
 *
 * <p>Contains extra logic for adding Tasks ({@link ResourceSlotMetricGroup}).
 */
@Internal
public class ResourceProviderConsumerMetricGroup extends ConsumerMetricGroup<ResourceProviderMetricGroup> {

    /** Map from attempt ID (task identifier) to task metrics. */
    private final Map<ConsumeAttemptID, ResourceSlotMetricGroup> slots = new HashMap<>();

    // ------------------------------------------------------------------------

    public ResourceProviderConsumerMetricGroup(
            MetricRegistry registry,
            ResourceProviderMetricGroup parent,
            ConsumerID consumer,
            @Nullable String jobName) {
        super(
                registry,
                parent,
                consumer,
                jobName,
                registry.getScopeFormats()
                        .getTaskManagerJobFormat()
                        .formatScope(checkNotNull(parent), consumer, jobName));
    }

    public final ResourceProviderMetricGroup parent() {
        return parent;
    }

    // ------------------------------------------------------------------------
    //  adding / removing tasks
    // ------------------------------------------------------------------------

    public ResourceSlotMetricGroup addTask(
            final ConsumeAttemptID executionAttemptID, final String resourceSlotName) {
        checkNotNull(executionAttemptID);
        checkNotNull(resourceSlotName);

        synchronized (this) {
            if (!isClosed()) {
                ResourceSlotMetricGroup prior = slots.get(executionAttemptID);
                if (prior != null) {
                    return prior;
                } else {
                    ResourceSlotMetricGroup task =
                            new ResourceSlotMetricGroup(registry, this, executionAttemptID, resourceSlotName);
                    slots.put(executionAttemptID, task);
                    return task;
                }
            } else {
                return null;
            }
        }
    }

    public void removeTaskMetricGroup(ConsumeAttemptID executionId) {
        checkNotNull(executionId);

        // this can be a call from this.close which iterates over tasks
        // changing tasks here would break iteration
        synchronized (this) {
            if (!isClosed()) {
                slots.remove(executionId);
                // keep this group open even if tasks is empty - to re-use on new task submission
                // the group will be closed by TM with the release of the last job slot on this TM
            }
        }
    }

    // ------------------------------------------------------------------------
    //  Component Metric Group Specifics
    // ------------------------------------------------------------------------

    @Override
    protected Iterable<? extends ComponentMetricGroup> subComponents() {
        checkState(holdsLock(this));
        return slots.values();
    }
}
